
#include <string.h>
#include "stm32f4xx_hal.h"

#define STM32_FLASH_SIZE (1024 * 1024) // 1MB

static const uint32_t sector_addr_begin[] = {
    0x08000000, // Sector 0, 16 KB
    0x08004000, // Sector 1, 16 KB
    0x08008000, // Sector 2, 16 KB
    0x0800C000, // Sector 3, 16 KB
    0x08010000, // Sector 4, 64 KB
    0x08020000, // Sector 5, 128 KB
    0x08040000, // Sector 6, 128 KB
    0x08060000, // Sector 7, 128 KB
    0x08080000, // Sector 8, 128 KB
    0x080A0000, // Sector 9, 128 KB
    0x080C0000, // Sector 10,128 KB
    0x080E0000, // Sector 11,128 KB
};

int stm32flash_get_sector(uint32_t addr)
{
    int sector = -1;

    if (addr < sector_addr_begin[0] || addr >= (sector_addr_begin[0] + STM32_FLASH_SIZE)) {
        return sector; // Invalid address
    }

    for (uint32_t i = 0; i < sizeof(sector_addr_begin) / sizeof(sector_addr_begin[0]); i++) {
        if (addr >= sector_addr_begin[i]) {
            sector = i;
        } else {
            break;
        }
    }
    return sector;
}

int stm32flash_check_align(uint32_t addr)
{
    for (uint32_t i = 0; i < sizeof(sector_addr_begin) / sizeof(sector_addr_begin[0]); i++) {
        if (addr == sector_addr_begin[i]) {
            return 1;
        }
    }
    return 0;
}

int stm32flash_erase(uint32_t addr, uint32_t length)
{
    int ret = 0;
    int first_sector = stm32flash_get_sector(addr);
    int last_sector = stm32flash_get_sector(addr + length - 1);
    if (first_sector == -1 || last_sector == -1 || last_sector < first_sector) {
        return -1; // Invalid sector
    }

    FLASH_EraseInitTypeDef EraseInitStruct;
    EraseInitStruct.TypeErase = FLASH_TYPEERASE_SECTORS;
    EraseInitStruct.VoltageRange = FLASH_VOLTAGE_RANGE_3;
    EraseInitStruct.Sector = first_sector;
    EraseInitStruct.NbSectors = last_sector - first_sector + 1;

    HAL_FLASH_Unlock();
    uint32_t sector_error = 0;
    if (HAL_FLASHEx_Erase(&EraseInitStruct, &sector_error) != HAL_OK) {
        ret = -1;
    } else {
        ret = 0;
    }
    __HAL_FLASH_DATA_CACHE_DISABLE();
    __HAL_FLASH_INSTRUCTION_CACHE_DISABLE();
    __HAL_FLASH_DATA_CACHE_RESET();
    __HAL_FLASH_INSTRUCTION_CACHE_RESET();
    __HAL_FLASH_INSTRUCTION_CACHE_ENABLE();
    __HAL_FLASH_DATA_CACHE_ENABLE();
    HAL_FLASH_Lock();

    return ret;
}

int stm32flash_write(uint32_t addr, const uint8_t* data, uint32_t len)
{
    uint32_t index, head_fill_len, last_fill_len, program_word_len, program_word;
    const uint8_t *dst_data, *src_data;

    dst_data = (const uint8_t *)addr;
    src_data = data;

    head_fill_len = addr & 0x3;
    head_fill_len = 4 - head_fill_len;
    head_fill_len &= 0x3;

    if (len <= head_fill_len) {
        head_fill_len = len;
        program_word_len = 0;
        last_fill_len = 0;
    } else {
        program_word_len = (len - head_fill_len )/4;
        last_fill_len = (len - head_fill_len ) & 0x3;
    }

    HAL_FLASH_Unlock();
    for(index = 0; index < head_fill_len; index ++)
    {
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE,addr,*data);
        addr ++;
        data ++;
    }

    for(index = 0; index < program_word_len; index ++)
    {
        memcpy(&program_word,data,4);
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD,addr,program_word);
        addr += 4;
        data += 4;
    }

    for(index = 0; index < last_fill_len; index ++)
    {
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE,addr,*data);
        addr ++;
        data ++;
    }

    HAL_FLASH_Lock();

    for(index = 0; index < len; index ++)
    {
        if( dst_data[index] != src_data[index] )
            return -1;
    }

    return 0;
}
